package com.ssl.security.util;

import com.ssl.security.constants.KeyStoreTypeEnum;
import com.ssl.security.exception.SecurityException;
import org.bouncycastle.asn1.DERBMPString;
import org.bouncycastle.asn1.pkcs.PKCSObjectIdentifiers;
import org.bouncycastle.cert.jcajce.JcaX509ExtensionUtils;
import org.bouncycastle.crypto.engines.DESedeEngine;
import org.bouncycastle.crypto.engines.RC2Engine;
import org.bouncycastle.crypto.modes.CBCBlockCipher;
import org.bouncycastle.pkcs.PKCS12PfxPduBuilder;
import org.bouncycastle.pkcs.PKCS12SafeBag;
import org.bouncycastle.pkcs.PKCS12SafeBagBuilder;
import org.bouncycastle.pkcs.PKCSException;
import org.bouncycastle.pkcs.bc.BcPKCS12MacCalculatorBuilder;
import org.bouncycastle.pkcs.bc.BcPKCS12PBEOutputEncryptorBuilder;
import org.bouncycastle.pkcs.jcajce.JcaPKCS12SafeBagBuilder;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.security.*;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.Enumeration;

/**
 * 描述：pfx工具类
 *
 * @author ssl
 * @create 2021/3/19 0019.
 */
public class PKCS12Utils {


    /**
     * 获取PFX证书中的私钥
     *
     * @param keyStoreType keyStore类型 {@link KeyStoreTypeEnum#pkcs12}
     * @param pfxData      pfx证书二进制流
     * @param password     pfx访问口令
     * @return
     * @throws Exception
     */
    public static PrivateKey getPrivateKey(KeyStoreTypeEnum keyStoreType, byte[] pfxData, String password) throws Exception {
        PrivateKey privateKey = null;
        KeyStore keystore = getKeyStore(keyStoreType, pfxData, password);
        Enumeration<String> enums = keystore.aliases();
        String keyAlias;
        while (enums.hasMoreElements()) {
            keyAlias = enums.nextElement();
            if (keystore.isKeyEntry(keyAlias)) {
                privateKey = (PrivateKey) keystore.getKey(keyAlias, password.toCharArray());
            }
        }
        return privateKey;
    }


    /**
     * 获取PFX证书中的公钥证书
     *
     * @param keyStoreType keyStore类型 {@link KeyStoreTypeEnum#pkcs12}
     * @param pfxData      pfx证书二进制流
     * @param password     pfx访问口令
     * @return
     * @throws Exception
     */
    public static X509Certificate getX509Certificate(KeyStoreTypeEnum keyStoreType, byte[] pfxData, String password) throws Exception {
        X509Certificate x509Certificate = null;
        KeyStore keystore = getKeyStore(keyStoreType, pfxData, password);
        Enumeration<String> enums = keystore.aliases();
        String keyAlias;
        while (enums.hasMoreElements()) {
            keyAlias = enums.nextElement();
            if (keystore.isKeyEntry(keyAlias)) {
                x509Certificate = (X509Certificate) keystore.getCertificate(keyAlias);
            }
        }
        return x509Certificate;
    }


    /**
     * 合成SM2证书的PFX
     *
     * @param privKey 私钥
     * @param pubKey 公钥
     * @param cert 公钥证书
     * @param passwd 访问口令
     * @return
     * @throws SecurityException
     */
    public static byte[] synthesisSm2(PrivateKey privKey, PublicKey pubKey, X509Certificate cert, String passwd) throws SecurityException {
        JcaX509ExtensionUtils extUtils;
        try {
            extUtils = new JcaX509ExtensionUtils();
        } catch (NoSuchAlgorithmException e) {
            throw new SecurityException("合成 SM2 PFX 证书失败。", e);
        }

        PKCS12SafeBagBuilder eeCertBagBuilder;
        try {
            eeCertBagBuilder = new JcaPKCS12SafeBagBuilder(cert);
        } catch (IOException e) {
            throw new SecurityException("合成 SM2 PFX 证书失败。", e);
        }
        eeCertBagBuilder.addBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_friendlyName,
                new DERBMPString(""));
        eeCertBagBuilder.addBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_localKeyId,
                extUtils.createSubjectKeyIdentifier(pubKey));

        char[] passwdChars = passwd.toCharArray();
        PKCS12SafeBagBuilder keyBagBuilder = new JcaPKCS12SafeBagBuilder(privKey,
                new BcPKCS12PBEOutputEncryptorBuilder(
                        PKCSObjectIdentifiers.pbeWithSHAAnd3_KeyTripleDES_CBC,
                        new CBCBlockCipher(new DESedeEngine())).build(passwdChars));
        keyBagBuilder.addBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_friendlyName,
                new DERBMPString(""));
        keyBagBuilder.addBagAttribute(PKCSObjectIdentifiers.pkcs_9_at_localKeyId,
                extUtils.createSubjectKeyIdentifier(pubKey));

        PKCS12PfxPduBuilder pfxPduBuilder = new PKCS12PfxPduBuilder();
        PKCS12SafeBag[] certs = new PKCS12SafeBag[1];
        certs[0] = eeCertBagBuilder.build();
        try {
            pfxPduBuilder.addEncryptedData(new BcPKCS12PBEOutputEncryptorBuilder(
                            PKCSObjectIdentifiers.pbeWithSHAAnd40BitRC2_CBC,
                            new CBCBlockCipher(new RC2Engine())).build(passwdChars),
                    certs);
            pfxPduBuilder.addData(keyBagBuilder.build());
            return pfxPduBuilder.build(new BcPKCS12MacCalculatorBuilder(), passwdChars).getEncoded();
        } catch (IOException | PKCSException e) {
            throw new SecurityException("合成 SM2 PFX 证书失败。", e);
        }
    }


    /**
     * 合成PFX证书
     *
     * @param keyStoreType    keyStore类型 {@link KeyStoreTypeEnum#pkcs12}
     * @param privateKey      私钥
     * @param password        证书访问密钥
     * @param x509Certificate 公钥证书
     * @return
     * @throws SecurityException
     */
    public static byte[] synthesis(KeyStoreTypeEnum keyStoreType, PrivateKey privateKey, char[] password, X509Certificate x509Certificate) throws
            SecurityException {
        Certificate[] chain = {x509Certificate};
        KeyStore keystore;
        try {
            keystore = KeyStore.getInstance(keyStoreType.getValue());
            keystore.load(null, password);
            keystore.setKeyEntry("", privateKey, password, chain);
            ByteArrayOutputStream bytesos = new ByteArrayOutputStream();
            keystore.store(bytesos, password);
            byte[] bytes = bytesos.toByteArray();
            return bytes;
        } catch (GeneralSecurityException | IOException e) {
            throw new SecurityException("合成 PFX 证书失败。", e);
        }
    }

    /**
     * 获取keyStore
     *
     * @param keyStoreType
     * @param pfxData
     * @param password
     * @return
     * @throws Exception
     */
    private static KeyStore getKeyStore(KeyStoreTypeEnum keyStoreType, byte[] pfxData, String password) throws Exception {
        KeyStore keystore = KeyStore.getInstance(keyStoreType.getValue());
        keystore.load(new ByteArrayInputStream(pfxData), password.toCharArray());
        return keystore;
    }

}
